package com.linjiahao.nettyGateway.loadBalancer;

import com.alibaba.fastjson.JSON;
import com.linjiahao.commonall.constant.AuthConstant;
import com.linjiahao.nettyGateway.cluster.ClusterManager;
import com.linjiahao.nettyGateway.cluster.ServerNode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.reactive.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.reactive.EmptyResponse;
import org.springframework.cloud.client.loadbalancer.reactive.Request;
import org.springframework.cloud.client.loadbalancer.reactive.Response;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.util.CollectionUtils;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Random;

/**
 * @Author Shinelon
 * @Date 2023/1/14 18:25
 * @Version 1.0
 */
@Slf4j
public class WebsocketLoadBalancer implements ReactorServiceInstanceLoadBalancer {

    private final String serviceId;

    private ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;

    private ClusterManager clusterManager;

    public WebsocketLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                      String serviceId,
                                      ClusterManager clusterManager) {
        this.serviceId = serviceId;
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
        this.clusterManager = clusterManager;
    }


    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        if (this.serviceInstanceListSupplierProvider == null) {
            log.warn("No servers available for service: " + this.serviceId);
        }
        ServiceInstanceListSupplier supplier = (ServiceInstanceListSupplier)this.serviceInstanceListSupplierProvider.getIfAvailable(NoopServiceInstanceListSupplier::new);
        return (supplier.get()).next().map(serviceInstances -> getInstanceResponse(serviceInstances,request));
    }

    private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances,Request request) {
        if (instances.isEmpty()) {
            log.warn("No servers available for service: " + this.serviceId);
            return new EmptyResponse();
        } else {
            List<String> token = ((WebsocketLoadBalancerRequest) request).getRequest().getHeaders().get(AuthConstant.SUB_PROTOCOL);
            ServiceInstance instance = null;
            if(CollectionUtils.isEmpty(token)){
                int pos = Math.abs((new Random()).nextInt(1000));
                instance = (ServiceInstance)instances.get(pos % instances.size());
            }else {
                ServerNode server = clusterManager.getServer(token.get(0));
                ServiceInstance serviceInstance =null;
                if(server !=null){
                    log.info("ws 请求进行负载均衡");
                    serviceInstance = instances.stream().filter(v -> server.getInstanceId().equals(v.getMetadata().get("nacos.instanceId"))).findFirst().orElse(null);
                    log.info("ws 负载均衡节点为：{}", JSON.toJSONString(serviceInstance));
                }
                if(serviceInstance == null){
                    int pos = Math.abs((new Random()).nextInt(1000));
                    instance = (ServiceInstance)instances.get(pos % instances.size());
                }else {
                    instance = serviceInstance;
                }
            }
            return new DefaultResponse(instance);
        }
    }
}
